#!/usr/bin/python3
# -*- coding: utf-8; tab-width: 4; indent-tabs-mode: t -*-

import os
import re
import time
import glob
import shutil
import socket
import pathlib
import subprocess
import libtorrent
import lxml.html
import urllib.parse
import urllib.request
import mirrors.plugin
from dataclasses import dataclass
from libtorrent import session as torrent_session
from libtorrent import torrent_info
from libtorrent import torrent_handle


class Main:

    def __init__(self):
        self.stdCategoryDict = {
            "roms": CategoryInfo(
                r"mame-roms-([0-9\.]+).torrent",                        # local file regex pattern
                r"MAME ([0-9\.]+) ROMs \(merged\)",                     # remote a tag regex pattern
                "mame-roms-%s.torrent",                                 # local file pattern
            ),
            "chds": CategoryInfo(
                r"mame-chds-([0-9\.]+).torrent",
                r"MAME ([0-9\.]+) CHDs \(merged\)",
                "mame-chds-%s.torrent",
            ),
            "software-list-roms": CategoryInfo(
                r"mame-software-list-roms-([0-9\.]+).torrent",
                r"MAME ([0-9\.]+) Software List ROMs \(merged\)",
                "mame-software-list-roms-%s.torrent",
            ),
            "software-list-chds": CategoryInfo(
                r"mame-software-list-chds-([0-9\.]+).torrent",
                r"MAME ([0-9\.]+) Software List CHDs \(merged\)",
                "mame-software-list-chds-%s.torrent",
            ),
            "extras": CategoryInfo(
                r"mame-extras-([0-9\.]+).torrent",
                r"MAME ([0-9\.]+) EXTRAs",
                "mame-extras-%s.torrent",
            ),
            "multimedia": CategoryInfo(
                r"mame-multimedia-([0-9\.]+).torrent",
                r"MAME ([0-9\.]+) Multimedia",
                "mame-multimedia-%s.torrent",
            ),
            "samples": CategoryInfo(
                r"mame-samples-([0-9\.]+).torrent",
                None,
                "mame-samples-%s.torrent",
            ),
        }

        self.cfg = mirrors.plugin.params["config"]
        self.dataDir = mirrors.plugin.params["storage-file"]["data-directory"]

        if "categories" in self.cfg:
            if len(set(self.cfg["categories"]) - set(self.stdCategoryDict.keys())) > 0:
                raise Exception("error")
            self.categories = self.cfg["categories"]
        else:
            self.categories = list(self.stdCategoryDict.keys())

        self.downloadingDir = os.path.join(self.dataDir, ".downloading")
        self.seedingDir = os.path.join(self.dataDir, ".seeding")

    def run(self):
        os.makedirs(self.downloadingDir, exist_ok=True)
        os.makedirs(self.seedingDir, exist_ok=True)

        self._removeRedundantFiles()

        self._fetchTorrents()

        recordDict = {}

        downloader = TorrentDownloader()
        try:
            # prepare tmpDir for each torrent file in self.downloadingDir, add torrent file into downloader
            for c in self.categories:
                validParamList = [self._getSeedingParamsByTorrentFn(x) for x in self._scanDirForCategory(self.seedingDir, c)]
                for torrentFn in self._scanDirForCategory(self.downloadingDir, c):
                    param = self._getDownloadingParamsByTorrentFn(torrentFn)
                    if not os.path.exists(param.resultDirFullfn):
                        if len(validParamList) > 0:
                            # copy data of old version to tmpDir, to save download time hugely
                            self._oldResultDir2NewResultDir(validParamList[-1].torrentFullfn, param.torrentFullfn, validParamList[-1].resultDirFullfn, param.resultDirFullfn)
                        else:
                            os.mkdir(param.resultDirFullfn)
                    downloader.addTorrent(torrentFn, param.torrentFullfn, None, param.resultDirFullfn)
                    recordDict[torrentFn] = c

            # download, process progress and completed files
            print("Start downloading %d torrents..." % (len(recordDict)))
            downloader.startDownload()
            while True:
                ret = downloader.pollCompleteTorrent()
                if ret is None:
                    break

                if ret[0]:
                    # some torrents are complete
                    for torrentFn in ret[1]:
                        downloadingParam = self._getDownloadingParamsByTorrentFn(torrentFn)
                        seedingParam = self._getSeedingParamsByTorrentFn(torrentFn, beforeOrAfterSwitch=True)
                        Util.forceMove(downloadingParam.torrentFullfn, seedingParam.torrentFullfn)
                        if os.path.exists(downloadingParam.torrentCfgFullfn):
                            Util.forceMove(downloadingParam.torrentCfgFullfn, seedingParam.torrentCfgFullfn)
                        Util.forceMove(downloadingParam.resultDirFullfn, seedingParam.resultDirFullfn)
                        os.symlink(seedingParam.symlinkTarget, seedingParam.symlinkFullfn)

                    print("Download Complete:")
                    for name in ret[1]:
                        print("    %s" % (name))
                else:
                    # no torrent is complete
                    print("Progress:")
                    for torrentFn, s in ret[1].items():
                        print('    %s (%.2f%%) %.2f MB/s peers: %d' % (torrentFn, s.progress * 100, s.download_rate / 1000000, s.num_peers))

                    # wait for 5 minute before checking status next time, won't affect torrent downloading becaure it is in other threads
                    time.sleep(5 * 60)
        finally:
            downloader.dispose()

        # process mame-roms-0.37b5 for mame2000
        pass

        # process mame-roms-0.78 for mame2003 and mame2003plus
        pass

        # process mame-roms-0.139 for mame2010
        pass

        # process mame-roms-0.160 for mame2015 and mess2015
        pass

        # do deduplication for all changed categories
        print("Start deduplication...")
        for c in set(recordDict.values()):
            targetDirList = [self._getSeedingParamsByTorrentFn(x).resultDirFullfn for x in self._scanDirForCategory(self.seedingDir, c)]
            self._deduplicateDirList(targetDirList)

        print("All done.")

    def _removeRedundantFiles(self):
        # process self.seedingDir and self.dataDir
        if True:
            # get valid files
            validFiles = []
            for fn in os.listdir(self.seedingDir):
                if not fn.endswith(".torrent"):
                    continue
                if not self._isValidMameTorrentFn(fn):
                    continue

                param = self._getSeedingParamsByTorrentFn(fn)
                if not os.path.isfile(param.torrentFullfn) or os.path.islink(param.torrentFullfn):
                    continue
                if os.path.exists(param.torrentCfgFullfn):
                    if not os.path.isfile(param.torrentCfgFullfn) or os.path.islink(param.torrentCfgFullfn):
                        continue
                if not os.path.isdir(param.resultDirFullfn) or os.path.islink(param.resultDirFullfn):
                    continue
                if not os.path.islink(param.symlinkFullfn) or os.readlink(param.symlinkFullfn) != param.symlinkTarget:
                    continue

                validFiles.append(param.torrentFullfn)
                if os.path.exists(param.torrentCfgFullfn):
                    validFiles.append(param.torrentCfgFullfn)
                validFiles.append(param.resultDirFullfn)
                validFiles.append(param.symlinkFullfn)

            # remove redundant files
            for fullfn in glob.glob(os.path.join(self.dataDir, "*")):
                if fullfn != self.seedingDir and fullfn != self.downloadingDir and fullfn not in validFiles:
                    Util.forceDelete(fullfn)
            for fullfn in glob.glob(os.path.join(self.seedingDir, "*")):
                if fullfn not in validFiles:
                    Util.forceDelete(fullfn)

        # process self.downloadingDir
        if True:
            # get valid files
            validFiles = []
            for fn in os.listdir(self.downloadingDir):
                if not fn.endswith(".torrent"):
                    continue
                if not self._isValidMameTorrentFn(fn):
                    continue

                param = self._getDownloadingParamsByTorrentFn(fn)
                if os.path.isfile(param.torrentFullfn) and not os.path.islink(param.torrentFullfn):
                    validFiles.append(param.torrentFullfn)
                    if os.path.isfile(param.torrentCfgFullfn) and not os.path.islink(param.torrentCfgFullfn):
                        validFiles.append(param.torrentCfgFullfn)
                    if os.path.isdir(param.resultDirFullfn) and not os.path.islink(param.resultDirFullfn):
                        validFiles.append(param.resultDirFullfn)

            # remove redundant files
            for fullfn in glob.glob(os.path.join(self.downloadingDir, "*")):
                if fullfn not in validFiles:
                    Util.forceDelete(fullfn)

    def _fetchTorrents(self):
        resp = urllib.request.urlopen("https://pleasuredome.github.io/pleasuredome/mame", timeout=60)
        print("Web page successfully fetched.")

        root = lxml.html.parse(resp)
        for aTag in root.xpath(".//a"):
            if aTag.get("href") is None:
                continue
            localFn = self._getTorrentFnByATag(aTag)
            if localFn is not None and localFn not in os.listdir(self.seedingDir) and localFn not in os.listdir(self.downloadingDir):
                param = self._getDownloadingParamsByTorrentFn(localFn)
                Util.magnet2torrent(aTag.get("href"), param.torrentFullfn)
                print("Torrent file created: %s" % (param.torrentFullfn))

    # def _fetchLocalTorrents(self):
    #     selfDir = os.path.dirname(os.path.realpath(__file__))
    #     for fn in os.listdir(selfDir):
    #         if self._isValidMameTorrentFn(fn) and fn not in os.listdir(self.seedingDir) and fn not in os.listdir(self.downloadingDir):
    #             fullfn = os.path.join(self.seedingDir, fn)
    #             shutil.copy(fn, fullfn)
    #             print("Torrent file created: %s" % (fullfn))

    def _isValidMameTorrentFn(self, fn):
        for c in self.categories:
            if re.fullmatch(self.stdCategoryDict[c].localFnRegex, fn) is not None:
                return True
        return False

    def _getTorrentFnByATag(self, aTag):
        for c in self.categories:
            if self.stdCategoryDict[c].remoteATagTextRegx is None:
                continue
            m = re.fullmatch(self.stdCategoryDict[c].remoteATagTextRegx, aTag.text)
            if m is None:
                continue
            if not aTag.get("href").startswith("magnet:"):
                continue
            return self.stdCategoryDict[c].localFnPattern % (m.group(1))
        return None

    def _getDownloadingParamsByTorrentFn(self, torrentFn):
        return DownloadingParams(
            os.path.join(self.downloadingDir, torrentFn),                                       # torrent file full path
            os.path.join(self.downloadingDir, torrentFn + ".cfg"),                              # torrent cfg file full path
            os.path.join(self.downloadingDir, torrentFn.replace(".torrent", "")),               # download dir full path
        )

    def _getSeedingParamsByTorrentFn(self, torrentFn, beforeOrAfterSwitch=False):
        name = torrentFn.replace(".torrent", "")
        if beforeOrAfterSwitch:
            ti = torrent_info(os.path.join(self.downloadingDir, torrentFn))
        else:
            ti = torrent_info(os.path.join(self.seedingDir, torrentFn))
        return SeedingParams(
            os.path.join(self.seedingDir, torrentFn),                                                  # target torrent file full path
            os.path.join(self.seedingDir, torrentFn + ".cfg"),                                         # target torrent cfg file full path
            os.path.join(self.seedingDir, name),                                                       # full path of directory for bittorrent
            os.path.join(self.dataDir, name),                                                          # full path of symlink for httpdir/ftp
            os.path.join(os.path.basename(self.seedingDir), name, torrent_info.name(ti)),              # target of symlink for httpdir/ftp
        )

    def _oldResultDir2NewResultDir(self, oldTorrentFullfn, newTorrentFullfn, oldResultDirFullfn, newResultDirFullfn):
        srcTi = torrent_info(oldTorrentFullfn)
        dstTi = torrent_info(newTorrentFullfn)

        assert not os.path.exists(newResultDirFullfn)
        subprocess.check_call(["cp", "-r", oldResultDirFullfn, newResultDirFullfn])         # shutil.copytree does not use reflink copy

        srcMainDir = os.path.join(newResultDirFullfn, torrent_info.name(srcTi))
        dstMainDir = os.path.join(newResultDirFullfn, torrent_info.name(dstTi))

        assert os.path.isdir(srcMainDir)
        os.rename(srcMainDir, dstMainDir)

    def _scanDirForCategory(self, curDir, category):
        ret = []
        for fn in sorted(os.listdir(curDir)):
            if re.fullmatch(self.stdCategoryDict[category].localFnRegex, fn) is not None:
                ret.append(fn)
        return ret

    def _scanSeedingDirForCategory(self, category):
        ret = {}
        for fn in sorted(os.listdir(self.seedingDir)):
            if re.fullmatch(self.stdCategoryDict[category].localFnRegex, fn) is not None:
                assert fn.endswith(".torrent")
                name = fn.replace(".torrent", "")
                ret[name] = (
                    os.path.join(self.seedingDir, fn),            # target torrent file
                    os.path.join(self.dataDir, name),             # target dir
                )
        return ret

    def _deduplicateDirList(self, dirList):
        # FIXME
        pass


@dataclass
class CategoryInfo:
    localFnRegex: str
    remoteATagTextRegx: str
    localFnPattern: str


@dataclass
class DownloadingParams:
    torrentFullfn: str
    torrentCfgFullfn: str
    resultDirFullfn: str


@dataclass
class SeedingParams:
    torrentFullfn: str
    torrentCfgFullfn: str
    resultDirFullfn: str
    symlinkFullfn: str
    symlinkTarget: str


class TorrentDownloader:

    def __init__(self):
        self._torrentDict = {}
        self._session = None

    def addTorrent(self, name, torrentFilePath, torrentCfgFilePath, tmpDir):
        assert self._session is None
        ti = torrent_info(torrentFilePath)
        self._torrentDict[torrent_info.name(ti)] = (torrentFilePath, torrentCfgFilePath, tmpDir, ti, name)

    def startDownload(self):
        assert self._session is None

        self._session = torrent_session()
        try:
            # change settings
            settings = self._session.get_settings()
            settings["active_downloads"] = -1
            settings["active_seeds"] = -1
            self._session.apply_settings(settings)

            # add torrents
            for torrentFilePath, torrentCfgFilePath, tmpDir, ti, name in self._torrentDict.values():
                self._session.add_torrent({
                    'ti': ti,
                    'save_path': tmpDir,
                })
        except BaseException:
            self._session = None        # FIXME: destroy session?
            raise

    def pollCompleteTorrent(self):
        assert self._session is not None

        if len(self._session.get_torrents()) == 0:
            return None

        completeList = []
        statusDict = {}

        for handle in self._session.get_torrents():
            torrentFilePath, torrentCfgFilePath, tmpDir, ti, name = self._torrentDict[handle.name()]
            s = handle.status()
            if not s.is_finished:
                statusDict[name] = s
            else:
                Util.deleteRedundantFilesFromTorrentResultDir(tmpDir, handle=handle)
                self._session.remove_torrent(handle)
                completeList.append(name)

        if len(completeList) > 0:
            return (True, completeList)
        else:
            return (False, statusDict)

    def dispose(self):
        self._session = None            # FIXME: destroy session?
        self._torrentDict = {}


class Util:

    @staticmethod
    def forceDelete(path):
        if os.path.islink(path):
            os.remove(path)
        elif os.path.isfile(path):
            os.remove(path)
        elif os.path.isdir(path):
            shutil.rmtree(path)
        elif os.path.lexists(path):
            os.remove(path)             # other type of file, such as device node
        else:
            pass                        # path does not exist, do nothing

    @staticmethod
    def forceMove(srcPath, dstPath):
        Util.forceDelete(dstPath)                # FIXME: needed?
        shutil.move(srcPath, dstPath)

    @staticmethod
    def deleteRedundantFilesFromTorrentResultDir(resultDir, handle=None, torrent_path=None, torrent_cfg_path=None):
        if handle is not None:
            assert torrent_path is None and torrent_cfg_path is None
            ti = torrent_handle.torrent_file(handle)
        else:
            assert torrent_path is not None
            ti = torrent_info(torrent_path)

        # remove files that are not in torrent
        pathList = [x.path for x in torrent_info.files(ti)]
        for dirpath, dirnames, filenames in os.walk(resultDir):
            for filename in filenames:
                fullfn = os.path.join(dirpath, filename)
                if Util.removePathPrefix(fullfn, resultDir) not in pathList:
                    os.remove(fullfn)

        # remove empty directores
        for dirpath, dirnames, filenames in os.walk(resultDir):
            for dirname in dirnames:
                fullfn = os.path.join(dirpath, dirname)
                if len(os.listdir(fullfn)) == 0:
                    os.rmdir(fullfn)

    @staticmethod
    def getFreeSocketPort(portType):
        if portType == "tcp":
            stlist = [socket.SOCK_STREAM]
        elif portType == "udp":
            stlist = [socket.SOCK_DGRAM]
        elif portType == "tcp+udp":
            stlist = [socket.SOCK_STREAM, socket.SOCK_DGRAM]
        else:
            assert False

        for port in range(10000, 65536):
            bFound = True
            for sType in stlist:
                s = socket.socket(socket.AF_INET, sType)
                try:
                    s.bind((('', port)))
                except socket.error:
                    bFound = False
                finally:
                    s.close()
            if bFound:
                return port

        raise Exception("no valid port")

    @staticmethod
    def removePathPrefix(fullfn, prefix):
        assert fullfn[:len(prefix + "/")] == prefix + "/"
        return fullfn[len(prefix + "/"):]

    @staticmethod
    def magnet2torrent(magnet_url, torrent_path=None):
        # Parse magnet URI parameters
        params = libtorrent.parse_magnet_uri(magnet_url)
        params.save_path = "/dev/null"                            # so that data won't be downloaded

        # Download torrent info
        session = libtorrent.session()
        handle = session.add_torrent(params)
        while not handle.status().has_metadata:
            time.sleep(1)

        # Download torrent data
        torrent_info = handle.status().torrent_file
        torrent_file = libtorrent.create_torrent(torrent_info)
        if torrent_path is None:
            torrent_path = torrent_info.name() + ".torrent"
        torrent_data = libtorrent.bencode(torrent_file.generate())

        # Save torrent file
        if os.path.isfile(torrent_path) and pathlib.Path(torrent_path).read_bytes() == torrent_data:
            # file already exists, avoid writing
            pass
        else:
            with open(torrent_path, "wb") as f:
                f.write(torrent_data)


###############################################################################

if __name__ == "__main__":
    Main().run()
